package cn.kevinking.ai.mcp.provider.heartbeat;

import cn.kevinking.ai.common.model.ServerInfo;
import cn.kevinking.ai.mcp.provider.McpProvider;
import cn.kevinking.ai.mcp.provider.service.NamingService;
import lombok.extern.slf4j.Slf4j;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

/**
 * @author kevinking
 */
@Slf4j
public class BeatManager {
    private static final String BEAT_THREAD_NAME = "mcp-provider-beat-sender";
    private final ScheduledThreadPoolExecutor scheduler;
    private final Map<String, BeatInfo> beatPool;
    private final NamingService serverProxy;

    public BeatManager(int threadCount, NamingService serverProxy) {
        this.scheduler = new ScheduledThreadPoolExecutor(threadCount, r -> {
            Thread thread = new Thread(r);
            thread.setDaemon(true);
            thread.setName(BEAT_THREAD_NAME);
            return thread;
        });
        this.beatPool = new ConcurrentHashMap<>(threadCount);
        this.serverProxy = serverProxy;
    }

    public BeatInfo newBeat(ServerInfo serverInfo, McpProvider instance) {
        return BeatInfo.create(serverInfo, instance);
    }

    public void addBeatInfo(BeatInfo beatInfo) {
        final String uniqueKey = buildBeatKey(beatInfo);
        BeatInfo existBeat;
        if ((existBeat = beatPool.remove(uniqueKey)) != null) {
            existBeat.setStatus(BeatStatus.STOPPED);
        }
        beatPool.put(uniqueKey, beatInfo);
        scheduler.schedule(new BeatTask(beatInfo, serverProxy), beatInfo.getPeriod(), TimeUnit.MILLISECONDS);
    }

    private String buildBeatKey(BeatInfo beatInfo) {
        return String.format("%s-%s:%s", beatInfo.getServerName(), beatInfo.getIp(), beatInfo.getPort());
    }
}
